#include <iostream>
#include <iomanip>
#include <string>
#include <vector>
#include "cartPoleSystem.h"
#include <casadi/casadi.hpp>
using namespace casadi;

int main() {
    CartPoleSystem cartPoleSystem;
    GuessStruct  guess  = cartPoleSystem.computeGuess();
    BoundsStruct bounds = cartPoleSystem.computeBounds(guess);
    // 
    int N = guess.control.size();
    int Num = 5*N + 2;
    
    //
    std::vector<MX> starts;
    std::vector<double> lbw;
    std::vector<double> ubw;
    std::vector<double> w0;

    starts.emplace_back(MX::sym("T0"));
    starts.emplace_back(MX::sym("Tf"));
    double WN = sqrt(1/9.81);
    double PI = 3.1415926;
    lbw.emplace_back(0);
    lbw.emplace_back(0.5*(2*PI*WN));
    ubw.emplace_back(0);
    ubw.emplace_back(4.0*(2*PI*WN));
    w0.emplace_back(0);
    w0.emplace_back(guess.time[guess.time.size()-1]);

    //
    MX J = MX::sym("J");
    J = 0;
    MX h = (starts[1] - starts[0])/(N-1);
    for(int i = 0; i < N; ++i)
    {
        std::string Name = "U"+std::to_string(i);
        MX u = MX::sym(Name);
        starts.emplace_back(u);
        lbw.emplace_back(bounds.control_lower);
        ubw.emplace_back(bounds.control_upper);
        w0.emplace_back(guess.control[i]);
        double fac = 1.0;
        if(i ==0 || i == N-1)
        {
            fac = 0.5;
        }
        J = J + u*u*fac*h;
    }
    
    //
    for(int j = 0; j < 4; ++j)
    {
        std::string Name = "X"+std::to_string(0)+"_"+std::to_string(j);
        starts.emplace_back(MX::sym(Name));
        lbw.emplace_back(bounds.initialState_lower(j));
        ubw.emplace_back(bounds.initialState_upper(j));
        Eigen::Matrix<double, 4,1> tmp = guess.state[0];
        w0.emplace_back(tmp(j));
    } 
    for(int i = 1; i < N-1; ++i)
    {
        for(int j = 0; j < 4; ++j)
        {
            std::string Name = "X"+std::to_string(i)+"_"+std::to_string(j);
            starts.emplace_back(MX::sym(Name));
            Eigen::Matrix<double, 4,1>  tmp = bounds.state_lower;
            lbw.emplace_back(tmp(j));
            tmp = bounds.state_upper;
            ubw.emplace_back(tmp(j));
            tmp = guess.state[i];
            w0.emplace_back(tmp(j));
        } 
    }
    for(int j = 0; j < 4; ++j)
    {
        std::string Name = "X"+std::to_string(N-1)+"_"+std::to_string(j);
        starts.emplace_back(MX::sym(Name));
        lbw.emplace_back(bounds.finalState_lower(j));
        ubw.emplace_back(bounds.finalState_upper(j));
        Eigen::Matrix<double, 4,1>  tmp = guess.state[N-1];
        w0.emplace_back(tmp(j));
    } 

    //
    std::vector<MX> g;
    std::vector<double> ubg;
    std::vector<double> lbg;
    for(int i = 0; i < N-1; ++i)
    {
        MX x0_0 = starts[2+N+4*i+0];
        MX x0_1 = starts[2+N+4*i+1];
        MX x0_2 = starts[2+N+4*i+2];
        MX x0_3 = starts[2+N+4*i+3];

        MX x1_0 = starts[2+N+4*i+0+4];
        MX x1_1 = starts[2+N+4*i+1+4];
        MX x1_2 = starts[2+N+4*i+2+4];
        MX x1_3 = starts[2+N+4*i+3+4];

        MX F1 = starts[2+i];
        MX F2 = starts[2+i+1];
        MX fx0_0 = cartPoleSystem.cartPoleDynamicsMX0(x0_0, x0_1, x0_2, x0_3, F1);
        MX fx0_1 = cartPoleSystem.cartPoleDynamicsMX1(x0_0, x0_1, x0_2, x0_3, F1);
        MX fx0_2 = cartPoleSystem.cartPoleDynamicsMX2(x0_0, x0_1, x0_2, x0_3, F1);
        MX fx0_3 = cartPoleSystem.cartPoleDynamicsMX3(x0_0, x0_1, x0_2, x0_3, F1);

        MX fx1_0 = cartPoleSystem.cartPoleDynamicsMX0(x1_0, x1_1, x1_2, x1_3, F2);
        MX fx1_1 = cartPoleSystem.cartPoleDynamicsMX1(x1_0, x1_1, x1_2, x1_3, F2);
        MX fx1_2 = cartPoleSystem.cartPoleDynamicsMX2(x1_0, x1_1, x1_2, x1_3, F2);
        MX fx1_3 = cartPoleSystem.cartPoleDynamicsMX3(x1_0, x1_1, x1_2, x1_3, F2);

        g.emplace_back((fx0_0+fx1_0)*h/2+x0_0-x1_0);
        g.emplace_back((fx0_1+fx1_1)*h/2+x0_1-x1_1);
        g.emplace_back((fx0_2+fx1_2)*h/2+x0_2-x1_2);
        g.emplace_back((fx0_3+fx1_3)*h/2+x0_3-x1_3);

        lbg.emplace_back(0);
        lbg.emplace_back(0);
        lbg.emplace_back(0);
        lbg.emplace_back(0);
        ubg.emplace_back(0);
        ubg.emplace_back(0);
        ubg.emplace_back(0);
        ubg.emplace_back(0);
    }
    for(int i = 0; i < Num; ++i)
    {
        std::cout<<starts[i]<<","<<w0[i]<<",["<<lbw[i]<<","<<ubw[i]<<"]"<<std::endl;
    }
    MXDict nlp;                 // NLP declaration
    nlp["x"]= vertcat(starts);  // decision vars
    nlp["f"] = J;               // objective
    nlp["g"] = vertcat(g);      // constraints
    std::cout<<"objective:"<<std::endl;
    std::cout<<J<<std::endl;
    std::cout<<"constraints:"<<std::endl;
    for(int i = 0; i < g.size(); i++)
    {
        std::cout<<g[i]<<std::endl;
    }
    
    // Create solver instance
    Function F = nlpsol("F","ipopt",nlp);
    // Solve the problem using a guess
    DMDict res = F(DMDict{{"x0",DM(w0)},{"ubx",ubw},{"lbx", lbw},{"ubg",ubg},{"lbg",lbg}});
    std::cout<<res["x"]<<std::endl;
    return 0;
}